#!/usr/bin/env python

import sys
import socket
import time
import getopt
import threading
import signal
import errno
import os
import random as _random
from random import randint

from threading import RLock, Thread

_random.seed()

def random( maxvalue ):
    return _random.random() * maxvalue
# random()

if sys.platform == "win32":
    timer = time.clock
else:
    timer = time.time


def debug( level, * msgs ):
    if debug.level >= level:
        debug.lock.acquire()
        try:
            print >> sys.stderr, threading.currentThread().getName(), ":",
            for m in msgs:
                print >> sys.stderr, m,
            print >> sys.stderr
        finally:
            debug.lock.release()
# debug()
debug.level = 0
debug.lock = RLock()


class Performer( object ):
    rid_lock = RLock()
    request_id = 0

    def __init__( self, server_address ):
        """Creates a new Performer that will be connected to server address.

        @param server_address: a tuple C{( address, port )} to connect to:
                               C{( 'localhost', 8000 )}
        @type  server_address: L{tuple}
        """
        self.socket = socket.socket( socket.AF_INET, socket.SOCK_STREAM )
        try:
            self.socket.connect( server_address )
        except socket.error, e:
            raise SystemExit( "Could not connect to '%s:%s'" % server_address )

        self.rfile = self.socket.makefile( 'r' )
        self.wfile = self.socket.makefile( 'w' )
        self.done = 0
    # __init__()

    def next_request_id( classobj ):
        """Increment internal request ID.

        @note: the request_id is shared among all StressTest descendant, so
        it's safe to call multiple tests, the request_id will not be an already
        used one.

        @return: the new request_id.
        @rtype: L{int}
        """
        classobj.rid_lock.acquire()
        classobj.request_id += 1
        v = classobj.request_id
        classobj.rid_lock.release()
        return v
    # inc_counter()
    next_request_id = classmethod( next_request_id )


    def set_request_id( classobj, request_id ):
        """Set request_id to new value.

        @param request_id: new value
        @type  request_id: L{int}
        """
        if not isinstance( request_id, int ):
            raise TypeError( "request_id must be an integer!" )
        classobj.rid_lock.acquire()
        classobj.request_id = request_id
        classobj.rid_lock.release()
    # set_request_id()
    set_request_id = classmethod( set_request_id )


    def send( self, value ):
        """Send value over the socket.

        @param value: what to send over the socket.
        """
        count = 0
        maxcount = 3
        while count < maxcount:
            try:
                self.wfile.write( value )
                self.wfile.flush()
                return
            except socket.error, e:
                if e[ 0 ] != errno.EINTR:
                    debug( 1, "socket error: %s" % e )
                    debug( 1, "Trying again... %s/%s" % ( count, maxcount ) )
                count += 1
    # send()
    write = send


    def send_command( self, command, *arguments ):
        """Send a command over the socket.

        This will use the request_id and call
        L{next_request_id()<StressTest.next_request_id>} after it.

        @param command: command name
        @param arguments: arguments to command.
        """
        cmd = "%s %d %s\n" % ( command,
                               self.next_request_id(),
                               " ".join( [ str( a ) for a in arguments ] ) )
        debug( 4, "send:", repr( cmd ) )
        self.send( cmd )
    # send_command()


    def recv( self, bytes ):
        """Receive C{bytes} from socket.

        @return: read bytes from socket.
        @rtype: L{str}
        """
        count = 0
        maxcount = 3
        while count < maxcount:
            try:
                v = self.rfile.read( bytes )
                debug( 3, "recv:", repr( v ) )
                return v
            except socket.error, e:
                if e[ 0 ] != errno.EINTR:
                    debug( 1, "socket error: %s" % e )
                    debug( 1, "Trying again... %s/%s" % ( count, maxcount ) )
                count += 1
    # recv()
    read = recv


    def recv_line( self ):
        """Receive a line (with trailing \n).

        @return: the read line.
        @rtype: L{str}
        """
        v = self.rfile.readline()
        debug( 4, "recv:", repr( v ) )
        return v
    # recv_line()
    readline = recv_line


    def recv_reply( self ):
        """Receive reply to command.

        The protocol states that a reply is a bunch of lines and ends at '\n\n'
        so we return everything but these two new line.

        @return: reply to command
        @rtype: L{str}
        """
        lines = []
        while True:
            l = self.readline()
            lines.append( l[ : -1 ] ) # chomp \n
            if len( lines ) > 1 and l == '\n':
                break

        del lines[ -1 ] # remove the last empty line
        return "\n".join( lines )
    # recv_reply()


    def recv_ok( self ):
        """Receive the acknowledgment from server.

        @return: True if received the ack, False otherwise.
        @rtype: L{bool}
        """
        ok_str = "+OK"
        v = self.recv_reply()
        if v != ok_str:
            debug( 1, "Could not receive %r, got %r instead." % ( ok_str, v ) )
            return False
        else:
            return True
    # recv_ok()


    def consult_account( self, account ):
        """Consult how much money an account have.

        @param account: what account to operate on.
        @return: Account money
        @rtype: L{int}
        """
        self.send_command( "consult", account )
        assert( self.recv_ok() )
        return int( self.recv_reply() )
    # consult_account()


    def perform( self, accounts=(), steps=20, delay=None ):
        """Perform tests.

        Returns a tuple of lists, one with the desired values, the order
        is the same as given in the C{accounts} parameter, and the statitics,
        which follow::

           ( total_time, sleep_times,
             try_again_times, try_again_tries,
             operation_times, operation_number )


        @param accounts: a list of accounts to operate on.
        @param steps: number of operations to do.
        @param delay: maximum delay to introduce.

        @return: tuple of values and statistics.
        @rtype: L{tuple}
        """
        values = [ 0 ] * len( accounts )

        # Avoid global name lookup
        _sleep  = time.sleep
        _random = random
        if delay:
            def sleep():
                v = random( delay )
                _sleep( v )
                sleep.times += v
            # sleep()
            sleep.times = 0.0
        else:
            def sleep():
                pass
            # sleep()
            sleep.times = 0.0

        len_accounts = len( accounts )
        start = timer()
        try_again_times  = [ 0.0, 0.0, 0.0 ]
        try_again_tries  = [   0,   0,   0 ]
        operation_times  = [ 0.0, 0.0, 0.0 ]
        operation_number = [   0,   0,   0 ]
        for i in xrange( steps ):
            # create values to deposit
            for j, account in enumerate( accounts ):
                value = randint( 1, 100 )

                op = randint( 0, 2 )
                while True:
                    debug( 2, "Try to execute command:", op )
                    this_op_time = timer()
                    try:
                        if   op == 0:
                            self.send_command( "withdraw", account, value )
                            assert( self.recv_ok() )
                            values[ j ] -= value
                        elif op == 1:
                            self.send_command( "deposit", account, value )
                            assert( self.recv_ok() )
                            values[ j ] += value
                        elif op == 2:
                            if len_accounts < 2:
                                op = 1
                                continue # cannot transfer with just one acc

                            other_acc = randint( 0, len_accounts - 1 )
                            if other_acc == j:
                                if other_acc == 0:
                                    other_acc = 1
                                else:
                                    other_acc -= 1
                            self.send_command( "transfer", account,
                                               accounts[ other_acc ], value )
                            assert( self.recv_ok() )
                            values[ j ] -= value
                            values[ other_acc ] += value
                        else:
                            raise Exception( "Oops... Invalid command" )
                        # endif

                        operation_times[ op ] += timer() - this_op_time
                        operation_number[ op ] += 1

                        self.done += 1

                        break

                    except AssertionError:
                        debug( 1, "Command failed to execute. Try again" )
                        errmsg = self.recv_reply()
                        debug( 2, "Not ok: %s" % errmsg )
                        t = max( delay, 0.001 )
                        _sleep( t )
                        try_again_tries[ op ] += 1
                        try_again_times[ op ] += t
                        continue
                # end while

                sleep()
        # for i in xrange( steps )
        stats = ( timer() - start, sleep.times,
                  try_again_times, try_again_tries,
                  operation_times, operation_number )
        if debug.level >= 3:
            debug( 3, """\
%(classname)s statistics:
 * total time...........: %(total_time)ss
 * time elapsed sleeping: %(sleep_time)ss
 * time elapsed in operations:
    * withdraw: %(withdraw_time)ss (%(withdraw_number)s)
    * deposit.: %(deposit_time)ss (%(deposit_number)s)
    * transfer: %(transfer_time)ss (%(transfer_number)s)
 * time elapsed trying again operations:
    * withdraw: %(withdraw_try_time)ss (%(withdraw_tries)s)
    * deposit.: %(deposit_try_time)ss (%(deposit_tries)s)
    * transfer: %(transfer_try_time)ss (%(transfer_tries)s)
""" % \
                   { 'classname': self.__class__.__name__,
                     'total_time': stats[ 0 ],
                     'sleep_time': stats[ 1 ],
                     'withdraw_try_time': stats[ 2 ][ 0 ],
                     'deposit_try_time': stats[ 2 ][ 1 ],
                     'transfer_try_time': stats[ 2 ][ 2 ],
                     'withdraw_tries': stats[ 3 ][ 0 ],
                     'deposit_tries': stats[ 3 ][ 1 ],
                     'transfer_tries': stats[ 3 ][ 2 ],
                     'withdraw_time': stats[ 4 ][ 0 ],
                     'deposit_time': stats[ 4 ][ 1 ],
                     'transfer_time': stats[ 4 ][ 2 ],
                     'withdraw_number': stats[ 5 ][ 0 ],
                     'deposit_number': stats[ 5 ][ 1 ],
                     'transfer_number': stats[ 5 ][ 2 ],
                } )
        return ( values, stats )
    # perform()


    def check( self, accounts, values ):
        """Check if the accounts have values.

        @param accounts: which accounts to check.
        @type  accounts: L{list} of L{int}s

        @param values: desired values.
        @type  values: L{list} of L{int}s

        @return: True if all values match.
        @rtype: L{bool}
        """
        self.send_command( "mconsult", *accounts )
        assert( self.recv_ok() )
        reply = self.recv_reply()
        server_values = [ int( s ) for s in reply.split() ]

        debug( 1, "Want: %s" % values )
        debug( 1, "Got.: %s" % server_values )

        for l, s in zip( values, server_values ):
            if l != s:
                debug( 0, "Different values: Local=%s, Server=%s." % ( l, s ) )
                return False

        return True
    # check()
# Performer


class MultiplePerformer( object ):
    class PerformerThread( Thread ):
        """Internal Thread to handle performers.

        This thread calls Performer's L{perform()<Performer.perform> with
        arguments in MultiplePerformer, those passed to
        L{MultiplePerformer.perform}.

        @note: these threads will be named with the prefix'Performer-'.
        """
        __performed_id = 0

        def __init__( self, parent ):
            self.parent = parent
            self.values = None
            self.performer = Performer( self.parent.server_address )

            name = "Performer-%d" % self.next_performer_id()
            Thread.__init__( self, name=name, target=self.action )
        # __init__()

        def action( self ):
            v, s = self.performer.perform( self.parent.accounts,
                                           self.parent.steps,
                                           self.parent.delay )
            self.values = v
            self.stats  = s
        # action()


        def next_performer_id( classobj ):
            """Increment performer id and returns it."""
            classobj.__performed_id += 1
            return classobj.__performed_id
        # next_performer_id()
        next_performer_id = classmethod( next_performer_id )
    # PerformerThread


    def __init__( self, server_address, number=20 ):
        """Creates a new Multiple Performer.


        @param server_address: a tuple C{( address, port )} to connect to:
                               C{( 'localhost', 8000 )}
        @type  server_address: L{tuple}

        @param number: number of threads to run simultaneously.
        @type  number: L{int}
        """
        self.server_address = server_address
        self.performers = []
        for i in xrange( number ):
            t = MultiplePerformer.PerformerThread( self )
            self.performers.append( t )
    # __init__()


    def perform( self, accounts, steps=20, delay=None ):
        """Perform tests.

        @param accounts: a list of accounts to operate on.
        @param steps: number of operations to do.
        @param delay: maximum delay to introduce.
        """
        self.accounts = accounts
        self.steps    = steps
        self.delay    = delay
        self.values   = [ 0 ] * len( accounts )

        self.base_vals = [ 0 ] * len( accounts )
        p = Performer( self.server_address )
        for i, account in enumerate( accounts ):
            self.base_vals[ i ] += p.consult_account( account )

        for t in self.performers:
            t.start()

        values = self.values
        for t in self.performers:
            # don't block forever, it will block signals!
            while t.isAlive():
                t.join( 0.1 )

            for i, v in enumerate( t.values ):
                values[ i ] += v
        self.values = values
    # perform()


    def check( self ):
        """Check if tests performed all right.

        @rtype: L{bool}
        """
        p = Performer( self.server_address )
        return p.check( self.accounts,
                        map( lambda a, b: a+b, self.values, self.base_vals )
                        )
    # check()


    def stats( self ):
        """Print out statistics"""
        gstats = [ 0.0, 0.0,
                   [ 0.0, 0.0, 0.0 ],
                   [   0,   0,   0 ],
                   [ 0.0, 0.0, 0.0 ],
                   [   0,   0,   0 ],
                   ]
        for t in self.performers:
            stats = t.stats
            gstats[ 0 ] += stats[ 0 ]
            gstats[ 1 ] += stats[ 1 ]
            for j in xrange( 2, 6 ):
                for i, v in enumerate( stats[ j ] ):
                    gstats[ j ][ i ] += v

            print """\
%(classname)s statistics:
 * total time...........: %(total_time)ss
 * time elapsed sleeping: %(sleep_time)ss
 * time elapsed in operations:
    * withdraw: %(withdraw_time)ss (%(withdraw_number)s)
    * deposit.: %(deposit_time)ss (%(deposit_number)s)
    * transfer: %(transfer_time)ss (%(transfer_number)s)
 * time elapsed trying again operations:
    * withdraw: %(withdraw_try_time)ss (%(withdraw_tries)s)
    * deposit.: %(deposit_try_time)ss (%(deposit_tries)s)
    * transfer: %(transfer_try_time)ss (%(transfer_tries)s)
""" % \
        { 'classname': t.getName(),
          'total_time': stats[ 0 ],
          'sleep_time': stats[ 1 ],
          'withdraw_try_time': stats[ 2 ][ 0 ],
          'deposit_try_time': stats[ 2 ][ 1 ],
          'transfer_try_time': stats[ 2 ][ 2 ],
          'withdraw_tries': stats[ 3 ][ 0 ],
          'deposit_tries': stats[ 3 ][ 1 ],
          'transfer_tries': stats[ 3 ][ 2 ],
          'withdraw_time': stats[ 4 ][ 0 ],
          'deposit_time': stats[ 4 ][ 1 ],
          'transfer_time': stats[ 4 ][ 2 ],
          'withdraw_number': stats[ 5 ][ 0 ],
          'deposit_number': stats[ 5 ][ 1 ],
          'transfer_number': stats[ 5 ][ 2 ],
          }
        # end for t in self.performers

        print "-"*78
        print """\
Global statistics:
 * total time...........: %(total_time)ss
 * time elapsed sleeping: %(sleep_time)ss (%(sleep_perc)0.2f%%)
 * time elapsed in operations:
    * withdraw: %(withdraw_time)ss (%(withdraw_number)s)
    * deposit.: %(deposit_time)ss (%(deposit_number)s)
    * transfer: %(transfer_time)ss (%(transfer_number)s)
 * time elapsed trying again operations: %(total_try_time)ss (%(total_try_perc)0.2f%%)
    * withdraw: %(withdraw_try_time)ss (%(withdraw_tries)s)
    * deposit.: %(deposit_try_time)ss (%(deposit_tries)s)
    * transfer: %(transfer_try_time)ss (%(transfer_tries)s)
""" % \
    { 'total_time': gstats[ 0 ],
      'sleep_time': gstats[ 1 ],
      'sleep_perc': gstats[ 1 ] * 100.0 / gstats[ 0 ],
      'total_try_time': sum( gstats[ 2 ] ),
      'total_try_perc': sum( gstats[ 2 ] ) * 100.0 / gstats[ 0 ],
      'withdraw_try_time': gstats[ 2 ][ 0 ],
      'deposit_try_time': gstats[ 2 ][ 1 ],
      'transfer_try_time': gstats[ 2 ][ 2 ],
      'withdraw_tries': gstats[ 3 ][ 0 ],
      'deposit_tries': gstats[ 3 ][ 1 ],
      'transfer_tries': gstats[ 3 ][ 2 ],
      'withdraw_time': gstats[ 4 ][ 0 ],
      'deposit_time': gstats[ 4 ][ 1 ],
      'transfer_time': gstats[ 4 ][ 2 ],
      'withdraw_number': gstats[ 5 ][ 0 ],
      'deposit_number': gstats[ 5 ][ 1 ],
      'transfer_number': gstats[ 5 ][ 2 ],
      }
    # stats()
# MultiplePerformer





if __name__ == "__main__":
    host    = "localhost"
    port    = 8021
    number  = 10
    rid     = 0
    test    = None
    delay   = 0.0
    verbose = 0
    steps   = 10
    stats   = False


    def usage():
        print """\
Usage:

   %(progname)s [options] <acc1> <acc2> ... <accN>

Where options can be:
  -H, --host=HOST        the coordinator hostname or IP.  [default=%(defhost)s]
  -p, --port=PORT        the coordinator port at host.    [default=%(defport)s]
  -n, --number=NUMBER    how many threads to run.         [default=%(defnum)s]
  -s, --steps=STEPS      steps per thread.                [default=%(defstep)s]
  -d, --delay=MAXDELAY   maximum delay to introduce       [default=%(defdelay)s]
  -r, --requestid=RID    first request identifier to use  [default=%(defrid)s]
  -S, --statistics       print out statistics
  -v, --verbose=LEVEL    verbosity level.                 [default=%(defverb)s]
  -h, --help             this message.

Where <acc1> <acc2> ... <accN> is a list of accounts to operate on.
""" % { "progname": sys.argv[ 0 ],
	"defhost": host,
	"defport": port,
        "defverb": verbose,
        "defnum": number,
        "defstep": steps,
        "defdelay": delay,
        "defrid": rid,
        }
        raise SystemExit
    # usage()

    opts, args = getopt.getopt( sys.argv[ 1: ], "hm:n:H:p:v:r:d:s:S",
                                [ "help", "number=", "delay=",
                                  "host=", "port=", "verbose=",
                                  "requestid=", "steps=", "statistics" ] )
    for o, a in opts:
        if o in ( "-n", "--number" ):
            number = int( a )
        elif o in ( "-H", "--host" ):
            host = a
        elif o in ( "-p", "--port" ):
            port = int( a )
        elif o in ( "-r", "--requestid" ):
            rid = int( a )
        elif o in ( "-d", "--delay" ):
            delay = float( a )
        elif o in ( "-s", "--steps" ):
            steps = int( a )
        elif o in ( "-v", "--verbose" ):
            verbose = int( a )
        elif o in ( "-h", "--help" ):
            usage()
        elif o in ( "-S", "--statistics" ):
            stats = True

    debug.level = verbose

    if not args:
        print "Please provide some accounts to operate on."
        raise SystemExit

    accounts = [ int( s ) for s in args ]
    mp = MultiplePerformer( ( host, port ), number )

    start_time = timer()
    def show_status( signum, frame ):
        total_steps = len( accounts ) * steps
        total_done  = 0
        print "-" * 78
        print "Threads status:"
        for t in mp.performers:
            done = t.performer.done
            total_done += done
            print "  * %s: %s steps done (%0.2f%%)" % \
                  ( t.getName(), done, done * 100.0 / total_steps )
        print "Total done: %s steps (%0.2f%%) in %s seconds." % \
              ( total_done, total_done * 100.0 / ( total_steps * number ),
                timer() - start_time )
        print "-" * 78
    # show_status()
    signal.signal( signal.SIGUSR1, show_status )

    debug( 1, "Current PID is %s" % os.getpid() )
    debug( 2, "You can send 'kill -SIGUSR1 %s' to get progress info." % \
           os.getpid() )


    mp.perform( accounts, steps, delay )
    if stats:
        mp.stats()
    print "Result:", mp.check()
# endif
